import os

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

data_dir = 'D:/pythonRoom/llm_llama/data'
'''
阿里百炼地址
https://bailian.console.aliyun.com/?tab=app#/api-key
'''
# 1. 配置阿里云API（替换成您的API Key）
print('初始化模型对象')
llm = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)

# 2. 使用本地嵌入模型（节省成本）
print('使用本地嵌入模型（节省成本）')
embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)

# 3. 加载文档
print('加载文档')
documents = SimpleDirectoryReader(data_dir).load_data()

# 4. 构建索引
print('构建索引')
index = VectorStoreIndex.from_documents(
    documents,
    embed_model=embed_model
)

# 5. 创建查询引擎
print('创建查询引擎')
query_engine = index.as_query_engine(llm=llm)

# 6. 提问
print('提问')
response = query_engine.query("李大长是如何发家致富的?")
print("回答:", response)
